'''
Created on Sep 1, 2019

@author: yl
'''

from scipy import signal
import numpy as np

def plane_pattern(X, Y):
    '''
    2d Gaussian Window
    '''
    window_x = signal.gaussian(pix_num, std=pix_num/8)
    window_y = signal.gaussian(1024, std=pix_num/8)
    window_x, window_y = np.meshgrid(window_x, window_y)
    window_2d = window_x*window_y
    
    '''
    Basic Pattern
    '''
    diff_Z_p = D * 2 / (1-V_m_x**2-V_m_y**2) + L * 2 * ((V_m_x**2+V_m_y**2)-(V_r_x**2+V_r_y**2)) / (1-V_m_x**2-V_m_y**2) / (1-V_r_x**2-V_r_y**2)
    #     d_ref = X*2*V_r_x*(1-V_r_x**2-V_r_y**2)/(1-V_r_x**2)/(1+V_r_x**2+V_r_y**2) + Y*2*V_r_y*(1-V_r_x**2-V_r_y**2)/(1-V_r_y**2)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
    #     d_mea = X*2*V_m_x*(1-V_m_x**2-V_m_y**2)/(1-V_m_x**2)/(1+V_m_x**2+V_m_y**2) + Y*2*V_m_y*(1-V_m_x**2-V_m_y**2)/(1-V_m_y**2)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_mea = (X*2*V_m_x+Y*2*V_m_y)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_ref = (X*2*V_r_x+Y*2*V_r_y)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
    diff_phi = k * (diff_Z_p + (d_mea-d_ref))
    I_beat = I_0*(1+np.cos(diff_phi))
    
    '''
    Newton Ring moving
    '''
    R = 0.2
    center = [(0)*pix_size, (0)*pix_size]
    r = np.sqrt((X-center[0])**2 + (Y-center[1])**2)
    d = R - np.sqrt(R**2 - r**2)
    A = 2 * np.pi * (d+2*D*0) / Lamda
    I_r_move = np.sin(A)**2 * I_0/5
    '''
    Newton Ring fixed
    '''
    
    '''
    加窗、取整
    '''
    Cam_Pattern = I_beat + I_r_move
    Cam_Pattern = Cam_Pattern * window_2d
    Cam_Pattern = Cam_Pattern.astype(np.int)
    
    
    return Cam_Pattern

def w_z(w_0, Z_R, Z):
    return w_0*np.sqrt(1 + (Z/Z_R)**2)

def R_z(Z_R, Z):
    return Z*(1+(Z_R/Z)**2)

def Phi_Gouy(Z_R, Z):
    return np.arctan(Z/Z_R)    

def Gaussian_pattern(X, Y):
    '''
    Basic Pattern
    '''
    Z_p_r = M + N + L * (2/(1-V_r_x**2-V_r_y**2) - 1)
    Z_p_m = M + N + (L+D) * (2/(1-V_m_x**2-V_m_y**2) - 1)
    
    diff_Z_p = D * 2 / (1-V_m_x**2-V_m_y**2) + L * 2 * ((V_m_x**2+V_m_y**2)-(V_r_x**2+V_r_y**2)) / (1-V_m_x**2-V_m_y**2) / (1-V_r_x**2-V_r_y**2)
    
    d_mea = (X*2*V_m_x+Y*2*V_m_y)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_ref = (X*2*V_r_x+Y*2*V_r_y)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
    
    r_mea = np.sqrt( X**2 + Y**2 + (L+D)**2 - ((2*X*V_m_x+2*Y*V_m_y+(L+D)*(1-V_m_x**2-V_m_y**2))/(1+V_m_x**2+V_m_y**2))**2 )
    r_ref = np.sqrt( X**2 + Y**2 + (L)**2 - ((2*X*V_r_x+2*Y*V_r_y+L*(1-V_r_x**2-V_r_y**2))/(1+V_r_x**2+V_r_y**2))**2 )
    
    R_ref = R_z(Z_R, Z_p_r+d_ref)
    R_mea = R_z(Z_R, Z_p_m+d_mea)
    
    Phi_gouy_ref = Phi_Gouy(Z_R, Z_p_r+d_ref)
    Phi_gouy_mea = Phi_Gouy(Z_R, Z_p_m+d_mea)
    
    diff_phi = k * (diff_Z_p + (d_mea-d_ref) + r_mea**2/2/R_mea - r_ref**2/2/R_ref) + Phi_gouy_mea - Phi_gouy_ref
    A_beat = 0.5 * I_0 * w_0**2 / w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea) * np.exp((-X**2-Y**2)/w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea))
#     A_beat = 0.5 * I_0
    I_beat = A_beat*(1+np.cos(diff_phi))
    
    '''
    Newton Ring moving
    '''
    R = 2
    center = [(pix_num-1-800)*pix_size/2, (pix_num+800)*pix_size/2]
    r = np.sqrt((X-center[0])**2 + (Y-center[1])**2)
    d = R - np.sqrt(R**2 - r**2)
    A = 2 * np.pi * (d+2*D) / Lamda
    I_r_2 = np.sin(A)**2 * I_0/20
    
    '''
    取整
    '''
    Cam_Pattern = I_beat
    Cam_Pattern = Cam_Pattern.astype(np.int)

    return Cam_Pattern

'''
Parameters Define
'''
j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]
k = 2 * np.pi / Lamda

'''
Camera Define
'''
frame_total = 10
fs_cam = 1000
I_0 = 127
pix_size = 5.3e-6
pix_num = 1280
screen_diameter = pix_num * pix_size
dx = np.linspace((-pix_num/2+1)*pix_size, pix_num/2*pix_size, num=pix_num)
dy = np.linspace((-1024/2+1)*pix_size, 1024/2*pix_size, num=1024)
X, Y = np.meshgrid(dx, dy)

'''
Beam Parameters
'''
M, N = 0.1, 0.05
L, D = 0.15, 0
w_0 = 1.12e-3
Z_R = np.pi * w_0**2 / Lamda

V_r_x, V_r_y, V_r_z = 0.0000, 0.000, 1 # Reference
V_m_x, V_m_y, V_z_m = 0.001, 0.001, 1 # Measurement

'''
Data generating
'''
Frame_set = []
for i in range(frame_total):
    D += Lamda / frame_total
    Frame_pattern = plane_pattern(X, Y)
#     Frame_pattern = Gaussian_pattern(X, Y)
    '''
    Noise
    '''
    for j in ranNoise_PSDn(Frame_pattern)):
        line_noise = np.random.normal(2.43, 0.8, len(Frame_pattern[j])) 
        line_noise = line_noise.round().astype(int)
        Frame_pattern[j] = Frame_pattern[j] + line_noise
    Frame_set.append(Frame_pattern) 
